"""
https://www.tensorflow.org/guide/keras/train_and_evaluate
"""

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, losses, metrics, optimizers, Model
from python_ai.common.xcommon import *
import os
import sys
import numpy as np
import random

random.seed(777)
np.random.seed(777)
tf.random.set_seed(777)

# model
inputs = keras.Input(shape=(784,), name="digits")
x = layers.Dense(64, activation="relu", name="dense_1")(inputs)
x = layers.Dense(64, activation="relu", name="dense_2")(x)
outputs = layers.Dense(10, activation="softmax", name="predictions")(x)

model = Model(inputs=inputs, outputs=outputs)

# data
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# Preprocess the data (these are NumPy arrays)
x_train = x_train.reshape(60000, 784).astype("float32") / 255
x_test = x_test.reshape(10000, 784).astype("float32") / 255

y_train = y_train.astype("float32")
y_test = y_test.astype("float32")

# Reserve 10,000 samples for validation
x_val = x_train[-10000:]
y_val = y_train[-10000:]
x_train = x_train[:-10000]
y_train = y_train[:-10000]

# specify the training config
is_use_cache = True
alpha = 0.001
ver = 'v1.0'
filename = os.path.basename(__file__)
save_path = './_save/' + filename + '_' + ver + '/weights.tmp.dat'

model.compile(
    optimizer=optimizers.RMSprop(learning_rate=alpha),
    loss=losses.SparseCategoricalCrossentropy(),
    metrics=[metrics.SparseCategoricalAccuracy()]
)

model.summary()

# fit and train
sep('Fit and train')
if is_use_cache and os.path.exists(save_path + '.index'):
    model.load_weights(save_path)
    print('LOADED')
else:
    history = model.fit(
        x_train,
        y_train,
        batch_size=64,
        epochs=2,
        validation_data=(x_val, y_val)
    )
    print(history.history)
    model.save_weights(save_path)
    print('SAVED')

sep('Evaluate on test data')
results = model.evaluate(x_test, y_test, batch_size=128)
print(f'test loss, test acc: {results}')

n_test = 10
m_test = len(x_test)
idx_test = range(m_test)
idx_selected = random.sample(idx_test, k=n_test)
print('idx_selected', idx_selected)
x_sel = x_test[idx_selected]
y_sel = y_test[idx_selected]
pred = model.predict(x_sel)
pred = np.argmax(pred, axis=1)
compare_tbl = np.c_[y_sel, pred, np.isclose(y_sel, pred)]
print(compare_tbl)
